package com.nepu.es

import java.util

import com.alibaba.fastjson.JSON
import com.alibaba.fastjson.serializer.SerializeConfig
import org.apache.http.HttpHost
import org.elasticsearch.action.bulk.BulkRequest
import org.elasticsearch.action.delete.DeleteRequest
import org.elasticsearch.action.index.IndexRequest
import org.elasticsearch.action.search.{SearchRequest, SearchResponse}
import org.elasticsearch.action.update.UpdateRequest
import org.elasticsearch.client._
import org.elasticsearch.common.text.Text
import org.elasticsearch.common.xcontent.XContentType
import org.elasticsearch.index.query.{BoolQueryBuilder, MatchQueryBuilder, QueryBuilders, RangeQueryBuilder}
import org.elasticsearch.index.reindex.UpdateByQueryRequest
import org.elasticsearch.script.{Script, ScriptType}
import org.elasticsearch.search.SearchHit
import org.elasticsearch.search.aggregations.{AggregationBuilders, Aggregations, BucketOrder}
import org.elasticsearch.search.aggregations.bucket.terms.{ParsedStringTerms, Terms, TermsAggregationBuilder}
import org.elasticsearch.search.aggregations.metrics.{AvgAggregationBuilder, ParsedAvg}
import org.elasticsearch.search.builder.SearchSourceBuilder
import org.elasticsearch.search.fetch.subphase.highlight.{HighlightBuilder, HighlightField}
import org.elasticsearch.search.sort.SortOrder


/**
  *
  * 本列是进行es的测试的
  *
  * @author chenshuaijun
  * @create 2022-11-04 11:55
  */
object ElasticSearchTest {

  val client: RestHighLevelClient = create()

  def main(args: Array[String]): Unit = {
    //    singleDataIdempotentInsert()
    //    singleDataNoIdempotentInsert()
    //    batchWrite()
    //    accordingIdModify()
    //    accordingConditionModify()
    //    deleteDoc()
    //    accordingConditionQuery()
    accordingConditionAgg()
    close()


  }

  /**
    * 单条数据写入-->幂等的写入
    */
  def singleDataIdempotentInsert(): Unit = {
    val indexRequest = new IndexRequest()
    // 指定索引的名字
    indexRequest.index("move_test")
    // 准备数据
    val movie = Movie("1001", "斗破苍穹")
    val movieJson: String = JSON.toJSONString(movie, new SerializeConfig(true))
    // 写入数据
    indexRequest.source(movieJson, XContentType.JSON)
    // 指定id
    indexRequest.id(movie.id)
    client.index(indexRequest, RequestOptions.DEFAULT)
  }

  /**
    * 单条数据写入-->非幂等的写入
    */
  def singleDataNoIdempotentInsert(): Unit = {
    val indexRequest = new IndexRequest()
    // 指定索引的名字
    indexRequest.index("move_test")
    // 准备数据
    val movie = Movie("1001", "斗罗大陆")
    val movieJson: String = JSON.toJSONString(movie, new SerializeConfig(true))
    // 写入数据
    indexRequest.source(movieJson, XContentType.JSON)
    // 不指定id
    // indexRequest.id(movie.id)
    client.index(indexRequest, RequestOptions.DEFAULT)
  }

  /**
    * 实现批量的写入
    */
  def batchWrite(): Unit = {
    val bulkRequest = new BulkRequest()
    val movies = List(
      Movie("1002", "我是谁"),
      Movie("1003", "你是谁"),
      Movie("1004", "他是谁"),
      Movie("1005", "俺也不知道")
    )
    for (movie <- movies) {
      val indexRequest = new IndexRequest()
      // 指定索引的名字,所以我们进行批量插入的时候，可以指定不同的索引
      indexRequest.index("move_test")
      val movieJson: String = JSON.toJSONString(movie, new SerializeConfig(true))
      indexRequest.source(movieJson, XContentType.JSON)
      // 我们使用批量幂等写入，所以这里指定id
      indexRequest.id(movie.id)
      bulkRequest.add(indexRequest)
    }
    client.bulk(bulkRequest, RequestOptions.DEFAULT)
  }

  /**
    * 修改数据：
    * 单条数据的修改
    */
  def accordingIdModify(): Unit = {
    val updateRequest = new UpdateRequest("move_test", "1001")
    updateRequest.doc("name", "我爱你")
    client.update(updateRequest, RequestOptions.DEFAULT)
  }

  /**
    * 修改数据
    * 根据查询条件修改
    */
  def accordingConditionModify(): Unit = {
    val updateByQueryRequest = new UpdateByQueryRequest()
    updateByQueryRequest.indices("move_index")
    val queryBuilder: MatchQueryBuilder = QueryBuilders.matchQuery("name", "wa ha ha")
    val map = new util.HashMap[String, AnyRef]()
    map.put("newValue", "shuang wai wai")
    val script = new Script(
      ScriptType.INLINE,
      Script.DEFAULT_SCRIPT_LANG,
      "ctx._source['name'] = params.newValue"
      , map
    )
    updateByQueryRequest.setScript(script)
    updateByQueryRequest.setQuery(queryBuilder)
    client.updateByQuery(updateByQueryRequest, RequestOptions.DEFAULT)
  }

  /**
    * 删除数据:指定索引和key就可以了
    */
  def deleteDoc(): Unit = {
    val deleteRequest = new DeleteRequest("move_index", "gXfcQYQBm4RGI4HeMm_y")
    client.delete(deleteRequest, RequestOptions.DEFAULT)
  }

  /**
    * 根据条件进行查询
    */
  def accordingConditionQuery(): Unit = {
    // 以下的代码是对这个代码做翻译
    val dsl: String =
      """
        |# * search :
        |# *
        |# * 查询 doubanScore>=5.0 关键词搜索 red sea
        |# * 关键词高亮显示
        |# * 显示第一页，每页 20 条
        |# * 按 doubanScore 从大到小排序
        |
        |GET move_index/_search
        |{
        |  "query": {
        |    "bool": {
        |      "filter": [
        |        {
        |          "range": {
        |            "doubanScore": {
        |              "gte": 5
        |            }
        |          }
        |        }
        |      ]
        |      ,
        |      "must": [
        |        {
        |          "match": {
        |            "name":"wai wai"
        |          }
        |        }
        |      ]
        |    }
        |  }
        |  , "sort": [
        |    {
        |      "doubanScore": {
        |        "order": "desc"
        |      }
        |    }
        |  ]
        |  ,
        |  "from": 0
        |  ,
        |  "size": 2
        |  ,
        |  "highlight": {
        |    "fields": {
        |      "name":{}
        |    }
        |  }
        |}
        |
      """.stripMargin
    val searchRequest = new SearchRequest()
    searchRequest.indices("move_index")
    val searchSourceBuilder = new SearchSourceBuilder()
    val boolQueryBuilder: BoolQueryBuilder = QueryBuilders.boolQuery()
    val rangeQueryBuilder: RangeQueryBuilder = new RangeQueryBuilder("doubanScore").gte("5.0")
    boolQueryBuilder.filter(rangeQueryBuilder)
    searchSourceBuilder.query(boolQueryBuilder)
    val matchQueryBuilder = new MatchQueryBuilder("name", "wai wai")
    boolQueryBuilder.must(matchQueryBuilder)

    searchSourceBuilder.sort("doubanScore", SortOrder.DESC)

    searchSourceBuilder.from(0)
    searchSourceBuilder.size(2)

    val highlightBuilder: HighlightBuilder = new HighlightBuilder().field("name")
    searchSourceBuilder.highlighter(highlightBuilder)

    searchRequest.source(searchSourceBuilder)

    val searchResponse: SearchResponse = client.search(searchRequest, RequestOptions.DEFAULT)

    // 处理结果
    // 首先拿到总的数据的条数
    val totalDocNum: Long = searchResponse.getHits.getTotalHits.value
    // 拿到具体的数据
    val hits: Array[SearchHit] = searchResponse.getHits.getHits
    for (searchHit <- hits) {
      // 拿到数据
      val sourceAsMap: util.Map[String, AnyRef] = searchHit.getSourceAsMap
      // 拿到高亮
      val highlightFields: util.Map[String, HighlightField] = searchHit.getHighlightFields
      // 因为只有一条数据我们直接拿出来
      val field: HighlightField = highlightFields.get("name")
      val fragments: Array[Text] = field.getFragments
      val highTag: String = fragments(0).toString
      sourceAsMap.put("name", highTag)

      val result: String = JSON.toJSONString(sourceAsMap, new SerializeConfig(true))
      println(result)
    }
  }

  /**
    * 根据条件进行聚合
    *
    */
  def accordingConditionAgg(): Unit = {
    // 以下的代码是对这个DSL的翻译
    val dsl: String =
      """
        |查询每位演员参演的电影的平均分，倒叙排序
        |GET move_index/_search
        |{
        |  "aggs": {
        |    "actorNameGroup": {
        |      "terms": {
        |        "field": "actorList.name.keyword",
        |        "size": 10
        |        ,
        |        "order": {
        |          "doubanScoreAvg": "asc"
        |        }
        |      }
        |      ,
        |      "aggs": {
        |        "doubanScoreAvg": {
        |          "avg": {
        |            "field": "doubanScore"
        |          }
        |        }
        |      }
        |    }
        |  }
        |  ,
        |  "size": 0
        |}
      """.stripMargin
    val searchRequest = new SearchRequest()
    searchRequest.indices("move_index")
    val searchSourceBuilder = new SearchSourceBuilder()
    val termsAggregationBuilder: TermsAggregationBuilder = AggregationBuilders.
      terms("actorNameGroup").field("actorList.name.keyword").size(10).
      order(BucketOrder.aggregation("doubanScoreAvg", false))
    val avgAggregationBuilder: AvgAggregationBuilder = AggregationBuilders.avg("doubanScoreAvg").field("doubanScore")
    searchSourceBuilder.aggregation(termsAggregationBuilder)
    termsAggregationBuilder.subAggregation(avgAggregationBuilder)
    searchRequest.source(searchSourceBuilder)
    val searchResponse: SearchResponse = client.search(searchRequest, RequestOptions.DEFAULT)

    val actorNameGroupTerms: ParsedStringTerms = searchResponse.getAggregations.get[ParsedStringTerms]("actorNameGroup")
    val buckets: util.List[_ <: Terms.Bucket] = actorNameGroupTerms.getBuckets
    import scala.collection.JavaConverters._
    for (elem <- buckets.asScala) {
      val key: String = elem.getKey.toString
      val count: Long = elem.getDocCount
      val avgScore: Double = elem.getAggregations.get[ParsedAvg]("doubanScoreAvg").getValue
      println(s"$key 参演了 $count 部电影，豆瓣平均分$avgScore")
    }
  }

  /**
    * 创建es客户端连接
    *
    * @return 客户端连接
    */
  def create(): RestHighLevelClient = {
    val builder: RestClientBuilder = RestClient.builder(new HttpHost("hadoop102", 9200))
    val highLevelClient = new RestHighLevelClient(builder)
    highLevelClient
  }

  /**
    * 关闭客户端连接：这个是必须的，要不然程序会一直被阻塞
    */
  def close(): Unit = {
    client.close()
  }

  case class Movie(id: String, name: String)

}
